﻿#include "scan_matching/nearest_search.h"
#include "mapping/probability_grid.h"

// #include <opencv2/opencv.hpp>

namespace cartographer {
namespace scan_matching {

static float dist_threshold = 3.0f;

/************** NearestSearch **************/
NearestSearch::NearestSearch() {
  dist_map_ptr_ = NULL;
}

NearestSearch::NearestSearch(const Grid2D& grid) {
    dist_map_ptr_ = NULL;
}

NearestSearch::~NearestSearch() {
  if(dist_map_ptr_ != NULL) {
    if(dist_map_ptr_->Cells != NULL) {
      free(dist_map_ptr_->Cells);
      dist_map_ptr_->Cells = NULL;
    }
    free(dist_map_ptr_);
    dist_map_ptr_ = NULL;
  }
}

void NearestSearch::BuildMapCSpaceByMaxProbability(const sensor::PointCloud& cloud, const transform::Rigid2d& initial_pose_estimate, const Grid2D& grid, const bool& is_use_best_slot) {
  if(dist_map_ptr_ != NULL) {
    if(dist_map_ptr_->Cells != NULL) {
      free(dist_map_ptr_->Cells);
      dist_map_ptr_->Cells = NULL;
    }
    free(dist_map_ptr_);
    dist_map_ptr_ = NULL;
  }
  std::vector<Eigen::Array2i> search_cell_index = GetMaxProbabilityPoint(cloud, initial_pose_estimate, grid, is_use_best_slot);
  dist_map_ptr_ = (DisMap*)malloc(sizeof(DisMap));
  dist_map_ptr_->OriginX = const_cast<Grid2D&>(grid).limits().max()[0];
  dist_map_ptr_->OriginY = const_cast<Grid2D&>(grid).limits().max()[1];
  dist_map_ptr_->SizeX = const_cast<Grid2D&>(grid).limits().cell_limits().num_x_cells; // judge by point cloud, cols
  dist_map_ptr_->SizeY = const_cast<Grid2D&>(grid).limits().cell_limits().num_y_cells; // rows
  dist_map_ptr_->Scale = const_cast<Grid2D&>(grid).limits().resolution();
  dist_map_ptr_->Cells = (MapCell*)NULL;

  const int nSize = sizeof(MapCell)*dist_map_ptr_->SizeX * dist_map_ptr_->SizeY;
  dist_map_ptr_->Cells = (MapCell*)malloc(nSize);

  memset(dist_map_ptr_->Cells, -1, nSize);
  //int size = search_cell_index.size();

  //std::cout << "size----:" << search_cell_index.size() << std::endl;
  for (int i = 0; i < search_cell_index.size(); i++) {

    Eigen::Array2i cell_idx = search_cell_index[i];

    int r = cell_idx[0];
    int c = cell_idx[1];    
    dist_map_ptr_->Cells[MAP_INDEX(dist_map_ptr_, r, c)].State = 1;
  }

  MapUpdateCspace(dist_map_ptr_, dist_threshold);

#if 0
  // show likelihold
  cv::Mat likelihold = cv::Mat(dist_map_ptr_->SizeY, dist_map_ptr_->SizeX, CV_8UC1, cv::Scalar(255));

  for (int i = 0; i < dist_map_ptr_->SizeY; i++)
  {
    for (int j = 0; j < dist_map_ptr_->SizeX; j++)
    {
      if (dist_map_ptr_->Cells[i*dist_map_ptr_->SizeX + j].Dist < dist_threshold)
      {
        likelihold.at<uchar>(i, j) = dist_map_ptr_->Cells[i*dist_map_ptr_->SizeX + j].Dist;
      }

    }
  }

  cv::imshow("likelihold",likelihold);
  cv::waitKey(10);
#endif

}

std::vector<Eigen::Array2i> NearestSearch::GetMaxProbabilityPoint(const sensor::PointCloud& cloud, const transform::Rigid2d& initial_pose_estimate, const Grid2D& grid, const bool& is_use_best_slot) {
    sensor::PointCloud input_transformed;
    if(is_use_best_slot) {
      input_transformed = cloud;
    } else {
      transform::Rigid3d pose_estimate_3d = transform::Embed3D(initial_pose_estimate);
      // Perform the actual transformation computation
      // Apply guessed transformation prior to search for neighbours
      input_transformed = sensor::TransformPointCloud3D(cloud, pose_estimate_3d);
    }

    int idx_threshold = 5;
    
    std::cout << " input_transformed ======= : " << input_transformed.size() << std::endl;

    std::vector<Eigen::Array2i> search_cell_index;

    for(int i = 0; i < input_transformed.size(); i++) {
        Eigen::Vector3f predict_pt3d = input_transformed[i].position;
        Eigen::Vector2f predict_pt(predict_pt3d[0], predict_pt3d[1]);
        Eigen::Array2i cell_idx = grid.limits().GetCellIndex(predict_pt); 

        // std::cout << " cell_idx ======= : " << cell_idx.transpose() << std::endl;

        //std::cout << " predict_pt : " << predict_pt.transpose() << std::endl;
        float best_probability = -1.0f;
        Eigen::Array2i best_cell_idx;
        // TODO : more than one best cell
        int x_lowwer = (cell_idx[0] - idx_threshold + 1);
        if(x_lowwer < 0) {
          x_lowwer = 0;
        }
        int x_upper = (cell_idx[0] + idx_threshold);

        if(x_upper > grid.limits().cell_limits().num_x_cells) {
          x_upper = grid.limits().cell_limits().num_x_cells;
        }
        int y_lowwer = (cell_idx[1] - idx_threshold + 1);
        if(y_lowwer < 0) {
          y_lowwer = 0;
        }

        int y_upper = (cell_idx[1] + idx_threshold);

        if(y_upper > grid.limits().cell_limits().num_y_cells) {
          y_upper = grid.limits().cell_limits().num_y_cells;
        }

        for(int x = x_lowwer; x < x_upper; x++) {
          for(int y = y_lowwer; y < y_upper; y++) {
              Eigen::Array2i search_cell_idx(x, y);
              if(!grid.limits().Contains(search_cell_idx)) {
                  continue;
              }

              float probability = static_cast<const ProbabilityGrid&>(grid).GetProbability(search_cell_idx);
              if(probability <= 0.5f * static_cast<const ProbabilityGrid&>(grid).GetMaxCorrespondenceCost()) {
                continue;
              }
              //std::cout << " probability ======= : " << probability << std::endl;
              //std::cout << " search_cell_idx ======= : " << x << " , " << y << " , " << search_cell_idx.transpose() << std::endl;

              if(probability >= best_probability) {
                 best_probability = probability;
                 best_cell_idx = search_cell_idx;
              }
          }         
        } 

        // std::cout << " best_probability ======= : " << best_probability << std::endl;

        if(best_probability >= 0.01f) {
            search_cell_index.push_back(best_cell_idx);
        }
    }

    // std::cout << " search_cell_index : " << search_cell_index.size() << std::endl;
    // exit(-1);

    return search_cell_index;
}

float NearestSearch::ComputeProbBilinearInterpolate(const Eigen::Array2i& cell_idx) const {
    float probability = 0.0f;

    int row = cell_idx[0];
    int col = cell_idx[1];
      
    if ((dist_map_ptr_->SizeX <= col) || (dist_map_ptr_->SizeY <= row) || (row < 0) || (col < 0)) {
        return probability;
    }

    // std::cout << "row:" << row << std::endl;
    // std::cout << "col:" << col << std::endl;
    // std::cout << "Map->SizeX:" << dist_map_ptr_->SizeX << std::endl;
    // std::cout << "Map->SizeY:" << dist_map_ptr_->SizeY << std::endl;

    // std::cout << "MAP_INDEX(dist_map_ptr_, col, row):" << MAP_INDEX(dist_map_ptr_, row, col) << std::endl;
    
    const int nSize = sizeof(MapCell)*dist_map_ptr_->SizeX * dist_map_ptr_->SizeY;

    if (MAP_INDEX(dist_map_ptr_, row, col) < 0 || MAP_INDEX(dist_map_ptr_, row, col) >= nSize) {
        return probability;
    }

    int origin_row = dist_map_ptr_->Cells[MAP_INDEX(dist_map_ptr_, row, col)].SrcJ;
    int origin_col = dist_map_ptr_->Cells[MAP_INDEX(dist_map_ptr_, row, col)].SrcI;
    
    // std::cout << " origin_row: " << origin_row << std::endl;
    // std::cout << " origin_col: " << origin_col << std::endl;

    if (origin_row < 0 || origin_col < 0) {
        return probability;
    }

    float distance = dist_map_ptr_->Cells[MAP_INDEX(dist_map_ptr_, row, col)].Dist;

    // std::cout << " distance 00: " << distance << std::endl;

    if (distance > dist_threshold) {
        return probability;
    }

    // std::cout << " distance : " << distance << std::endl;

    probability = ComputeProbabilityByGaussian(distance);

    return probability;
}

float NearestSearch::ComputeProbabilityByGaussian(const float& distance) const {
    float sigma = 1.5;
    float coff_part_one = (1.0f / (std::sqrt(2.0f * M_PI) * sigma));
    float coff_part_two = -2.0f * sigma * sigma;
    //return coff_part_one * exp(distance * distance / coff_part_two);
    return exp(distance * distance / coff_part_two);
}

} // namespace scan_matching
}
